// Copyright(C) 2023 InfiniFlow, Inc. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

module infinity_core:physical_insert.impl;

import :physical_insert;
import :query_context;
import :table_def;
import :data_table;
import :physical_operator_type;
import :operator_state;
import :expression_state;
import :data_block;
import :expression_evaluator;
import :base_expression;
import :default_values;
import :status;
import :infinity_exception;
import :logger;
import :meta_info;
import :txn_state;
import :wal_manager;
import :infinity_context;
import :new_txn;
import :cast_function;
import :bound_cast_func;
import :column_vector;
import :value;

import std;
import third_party;

import embedding_info;
import logical_type;
import data_type;
import column_def;

namespace infinity {

// Helper function to handle embedding dimension conversion
void HandleEmbeddingDimensionConversion(std::shared_ptr<ColumnVector> source_column,
                                        std::shared_ptr<ColumnVector> target_column,
                                        const EmbeddingInfo *source_info,
                                        const EmbeddingInfo *target_info) {
    size_t row_count = source_column->Size();
    i32 source_dim = source_info->Dimension();
    i32 target_dim = target_info->Dimension();

    if (source_dim == target_dim) {
        // Same dimensions, copy directly
        target_column->AppendWith(*source_column, 0, row_count);
        return;
    }

    // Handle dimension conversion
    for (size_t row_idx = 0; row_idx < row_count; ++row_idx) {
        Value source_value = source_column->GetValueByIndex(row_idx);
        std::span<char> source_data = source_value.GetEmbedding();

        if (source_info->Type() == EmbeddingDataType::kElemFloat) {
            const float *source_ptr = reinterpret_cast<const float *>(source_data.data());
            std::vector<float> target_vector(target_dim);

            if (source_dim <= target_dim) {
                // Pad with zeros if target dimension is larger
                std::copy(source_ptr, source_ptr + source_dim, target_vector.begin());
                std::fill(target_vector.begin() + source_dim, target_vector.end(), 0.0f);
            } else {
                // Truncate if target dimension is smaller
                std::copy(source_ptr, source_ptr + target_dim, target_vector.begin());
            }

            // Create target embedding value
            auto target_embedding_info = EmbeddingInfo::Make(EmbeddingDataType::kElemFloat, target_dim);
            Value target_value = Value::MakeEmbedding(reinterpret_cast<const char *>(target_vector.data()), target_embedding_info);
            target_column->AppendValue(target_value);
        } else if (source_info->Type() == EmbeddingDataType::kElemDouble) {
            const double *source_ptr = reinterpret_cast<const double *>(source_data.data());
            std::vector<double> target_vector(target_dim);

            if (source_dim <= target_dim) {
                // Pad with zeros if target dimension is larger
                std::copy(source_ptr, source_ptr + source_dim, target_vector.begin());
                std::fill(target_vector.begin() + source_dim, target_vector.end(), 0.0);
            } else {
                // Truncate if target dimension is smaller
                std::copy(source_ptr, source_ptr + target_dim, target_vector.begin());
            }

            // Create target embedding value
            auto target_embedding_info = EmbeddingInfo::Make(EmbeddingDataType::kElemDouble, target_dim);
            Value target_value = Value::MakeEmbedding(reinterpret_cast<const char *>(target_vector.data()), target_embedding_info);
            target_column->AppendValue(target_value);
        } else {
            UnrecoverableError(fmt::format("Unsupported embedding data type for dimension conversion"));
        }
    }
}

void PhysicalInsert::Init(QueryContext *query_context) {}

bool PhysicalInsert::Execute(QueryContext *query_context, OperatorState *operator_state) {
    StorageMode storage_mode = InfinityContext::instance().storage()->GetStorageMode();
    if (storage_mode == StorageMode::kUnInitialized) {
        UnrecoverableError("Uninitialized storage mode");
    }

    if (storage_mode != StorageMode::kWritable) {
        operator_state->status_ = Status::InvalidNodeRole("Attempt to write on non-writable node");
        operator_state->SetComplete();
        return true;
    }

    std::shared_ptr<DataBlock> output_block;

    // Check if we have a child operator (INSERT SELECT) or direct values (INSERT VALUES)
    if (operator_state->prev_op_state_ != nullptr) {
        // INSERT SELECT case: get data from child operator
        OperatorState *prev_op_state = operator_state->prev_op_state_;
        size_t data_block_count = prev_op_state->data_block_array_.size();

        if (data_block_count == 0) {
            // No data to insert
            std::unique_ptr<std::string> result_msg = std::make_unique<std::string>("INSERTED 0 Rows");
            if (operator_state == nullptr) {
                std::vector<std::shared_ptr<ColumnDef>> column_defs;
                std::shared_ptr<TableDef> result_table_def_ptr =
                    TableDef::Make(std::make_shared<std::string>("default_db"), std::make_shared<std::string>("Tables"), nullptr, column_defs);
                output_ = std::make_shared<DataTable>(result_table_def_ptr, TableType::kDataTable);
                output_->SetResultMsg(std::move(result_msg));
            } else {
                InsertOperatorState *insert_operator_state = static_cast<InsertOperatorState *>(operator_state);
                insert_operator_state->result_msg_ = std::move(result_msg);
            }
            operator_state->SetComplete();
            return true;
        }

        // Merge all data blocks from child operator
        output_block = DataBlock::Make();
        // Use target table column types for the output block
        std::vector<std::shared_ptr<DataType>> target_types;
        for (size_t i = 0; i < static_cast<size_t>(table_info_->column_count_); ++i) {
            target_types.emplace_back(table_info_->column_defs_[i]->type());
        }
        output_block->Init(target_types);

        // Check if we need to apply type casting for INSERT SELECT
        DataBlock *first_block = prev_op_state->data_block_array_[0].get();
        bool needs_casting = false;
        for (size_t i = 0; i < first_block->column_count() && i < target_types.size(); ++i) {
            if (*target_types[i] != *first_block->column_vectors[i]->data_type()) {
                needs_casting = true;
                break;
            }
        }

        if (needs_casting) {
            // Apply type casting for each block
            for (size_t block_idx = 0; block_idx < data_block_count; ++block_idx) {
                DataBlock *input_data_block_ptr = prev_op_state->data_block_array_[block_idx].get();

                // Cast each column if needed
                for (size_t col_idx = 0; col_idx < input_data_block_ptr->column_count() && col_idx < target_types.size(); ++col_idx) {
                    auto source_column = input_data_block_ptr->column_vectors[col_idx];
                    auto target_column = output_block->column_vectors[col_idx];
                    auto target_type = target_types[col_idx];

                    if (*source_column->data_type() == *target_type) {
                        // No casting needed, copy directly
                        target_column->AppendWith(*source_column, 0, source_column->Size());
                    } else if (source_column->data_type()->type() == LogicalType::kEmbedding && target_type->type() == LogicalType::kEmbedding) {
                        // Special handling for embedding dimension conversion
                        auto source_info = static_cast<const EmbeddingInfo *>(source_column->data_type()->type_info().get());
                        auto target_info = static_cast<const EmbeddingInfo *>(target_type->type_info().get());

                        if (source_info->Type() == target_info->Type()) {
                            // Same data type (float/double), handle dimension conversion
                            HandleEmbeddingDimensionConversion(source_column, target_column, source_info, target_info);
                        } else {
                            // Different data types, use standard casting
                            auto cast_func = CastFunction::GetBoundFunc(*source_column->data_type(), *target_type);
                            CastParameters cast_parameters;
                            cast_func.function(source_column, target_column, source_column->Size(), cast_parameters);
                        }
                    } else {
                        // Apply standard casting (this will handle BigInt->Integer, etc.)
                        auto cast_func = CastFunction::GetBoundFunc(*source_column->data_type(), *target_type);
                        CastParameters cast_parameters;
                        cast_func.function(source_column, target_column, source_column->Size(), cast_parameters);
                    }
                }
            }
        } else {
            // No casting needed, append directly
            for (size_t block_idx = 0; block_idx < data_block_count; ++block_idx) {
                DataBlock *input_data_block_ptr = prev_op_state->data_block_array_[block_idx].get();
                output_block->AppendWith(input_data_block_ptr);
            }
        }
        output_block->Finalize();

    } else {
        // INSERT VALUES case: evaluate expressions to create data block
        size_t row_count = value_list_.size();
        if (row_count == 0) {
            UnrecoverableError("No values to insert");
        }

        size_t column_count = value_list_[0].size();
        size_t table_collection_column_count = table_info_->column_count_;
        if (column_count != table_collection_column_count) {
            UnrecoverableError(
                fmt::format("Insert values count{} isn't matched with table column count{}.", column_count, table_collection_column_count));
        }

        // Prepare the output block
        std::vector<std::shared_ptr<DataType>> output_types;
        output_types.reserve(column_count);
        auto field_list = value_list_[0];
        size_t field_count = field_list.size();
        for (size_t i = 0; i < field_count; ++i) {
            auto data_type = field_list[i]->Type();
            output_types.emplace_back(std::make_shared<DataType>(data_type));
        }

        output_block = DataBlock::Make();
        output_block->Init(output_types);
        std::shared_ptr<DataBlock> output_block_tmp = DataBlock::Make();
        output_block_tmp->Init(output_types);

        ExpressionEvaluator evaluator;
        evaluator.Init(nullptr);
        // Each cell's expression of a column may differ. So we have to evaluate each cell instead of column here.
        for (size_t row_idx = 0; row_idx < row_count; ++row_idx) {
            // Reset the temporary block for each row
            output_block_tmp->Reset();

            for (size_t expr_idx = 0; expr_idx < column_count; ++expr_idx) {
                const std::shared_ptr<BaseExpression> &expr = value_list_[row_idx][expr_idx];
                std::shared_ptr<ExpressionState> expr_state = ExpressionState::CreateState(expr);
                evaluator.Execute(expr, expr_state, output_block_tmp->column_vectors[expr_idx]);
            }
            output_block->AppendWith(output_block_tmp);
        }
        output_block->Finalize();
    }

    NewTxn *new_txn = query_context->GetNewTxn();
    Status status = new_txn->Append(*table_info_, output_block);
    if (!status.ok()) {
        operator_state->status_ = status;
    }

    std::unique_ptr<std::string> result_msg = std::make_unique<std::string>(fmt::format("INSERTED {} Rows", output_block->row_count()));
    if (operator_state == nullptr) {
        // Generate the result table
        std::vector<std::shared_ptr<ColumnDef>> column_defs;
        std::shared_ptr<TableDef> result_table_def_ptr =
            TableDef::Make(std::make_shared<std::string>("default_db"), std::make_shared<std::string>("Tables"), nullptr, column_defs);
        output_ = std::make_shared<DataTable>(result_table_def_ptr, TableType::kDataTable);
        output_->SetResultMsg(std::move(result_msg));
    } else {
        InsertOperatorState *insert_operator_state = static_cast<InsertOperatorState *>(operator_state);
        insert_operator_state->result_msg_ = std::move(result_msg);
    }
    operator_state->SetComplete();
    return true;
}

} // namespace infinity
